#!/usr/bin/env python3

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

import warnings

{%- macro force_import(name) %}
import fbgemm_gpu.split_embedding_codegen_lookup_invokers.{{ name }} as {{ name }}  # noqa: F401
{%- endmacro %}

{%- macro try_import(name) %}
try:
    # Import is placed under a try-except bc the op is experimental and can be 
    # removed/updated in the future
    import fbgemm_gpu.split_embedding_codegen_lookup_invokers.{{ name }} as {{ name }}  # noqa: F401
except:
    warnings.warn(
        f"""\033[93m
        Failed to import: fbgemm_gpu.split_embedding_codegen_lookup_invokers.{{ name }}
        \033[0m""",
        DeprecationWarning,
    )
{%- endmacro %}

# TBE optimizers
{{- force_import("lookup_args") }}
{%- for optim in all_optimizers %}
    {{ try_import("lookup_" + optim) }}
{%- endfor %}

# SSD TBE optimizers
{{- try_import("lookup_args_ssd") }}
{%- for optim in ssd_optimizers %}
    {{ try_import("lookup_" + optim + "_ssd") }}
{%- endfor %}
